{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f348ee10",
   "metadata": {},
   "source": [
    "# 测试 `relay.qnn.simulated_quantize`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1d7a33ab",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import set_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edc71c5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "from tvm import te\n",
    "import numpy as np\n",
    "from tvm import relay\n",
    "from tvm.contrib import graph_executor\n",
    "from tvm.runtime.vm import VirtualMachine\n",
    "from tvm.topi.nn.qnn import SQNN_DTYPE_TO_CODE\n",
    "from utils import allclose_with_rounding, quantize_test_driver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b044c702",
   "metadata": {},
   "outputs": [],
   "source": [
    "def allclose_with_rounding(a, b):\n",
    "    \"\"\"\n",
    "    比较两个数组是否足够接近，允许一定数量的四舍五入误差。\n",
    "    用于QNN模拟量化测试中，验证量化前后结果的一致性。\n",
    "\n",
    "    参数:\n",
    "    a: 第一个比较数组 (NumPy数组或TVM张量)\n",
    "    b: 第二个比较数组 (NumPy数组或TVM张量)\n",
    "\n",
    "    异常:\n",
    "    AssertionError: 如果不匹配元素数量超过3个\n",
    "\n",
    "    说明:\n",
    "    - 函数通过逐元素比较创建布尔数组mismatch\n",
    "    - 使用np.sum计算不匹配元素总数\n",
    "    - 允许最多3个不匹配元素，以容忍GPU浮点运算的舍入误差\n",
    "    \"\"\"\n",
    "    # 创建布尔数组，True表示对应位置元素不相等\n",
    "    mismatch = a != b\n",
    "    # 计算不匹配元素总数并断言不超过3个\n",
    "    # 容忍少量误差是因为GPU的fp32算术可能有舍入误差\n",
    "    assert np.sum(mismatch) <= 3\n",
    "\n",
    "def quantize_test_driver(in_dtype, quant_args, axis, out_dtype, in_data):\n",
    "    \"\"\"\n",
    "    QNN模拟量化测试的驱动函数，用于执行量化操作并返回结果。\n",
    "\n",
    "    该函数构建一个包含QNN量化操作的Relay计算图，编译并执行它，\n",
    "    最后返回量化后的结果。主要用于测试模拟量化功能的正确性。\n",
    "\n",
    "    参数:\n",
    "    in_dtype (str): 输入数据的数据类型，如 'float32'\n",
    "    quant_args (dict): 量化参数字典，必须包含以下键:\n",
    "        - 'out_zero_point': 输出量化的零点值\n",
    "        - 'out_scale': 输出量化的缩放因子\n",
    "    axis (int): 量化操作应用的轴\n",
    "    out_dtype (str): 输出数据的数据类型，如 'int8', 'uint8'\n",
    "    in_data (numpy.ndarray): 输入数据数组\n",
    "\n",
    "    返回:\n",
    "    numpy.ndarray: 量化后的输出数据\n",
    "    \"\"\"\n",
    "    # 获取输入数据的形状\n",
    "    shape = in_data.shape\n",
    "    # 创建Relay变量表示输入数据\n",
    "    input_data = relay.var(\"input_data\", shape=shape, dtype=in_dtype)\n",
    "    # 创建输出零点和缩放因子的常量\n",
    "    output_zero_point = relay.const(quant_args[\"out_zero_point\"])\n",
    "    output_scale = relay.const(quant_args[\"out_scale\"])\n",
    "    # 执行QNN量化操作\n",
    "    quantized_output = relay.qnn.quantize(\n",
    "        input_data,\n",
    "        output_scale=output_scale,\n",
    "        output_zero_point=output_zero_point,\n",
    "        axis=axis,\n",
    "        out_dtype=out_dtype,\n",
    "    )\n",
    "    # 构建Relay函数和IRModule\n",
    "    mod = relay.Function(relay.analysis.free_vars(quantized_output), quantized_output)\n",
    "    mod = tvm.IRModule.from_expr(mod)\n",
    "    # 编译模型\n",
    "\n",
    "    dev = tvm.cpu(0)\n",
    "    target = \"llvm\"\n",
    "    with tvm.transform.PassContext(opt_level=3):\n",
    "        lib = relay.build(mod, target=target)\n",
    "    # 创建运行时模块\n",
    "    rt_mod = graph_executor.GraphModule(lib[\"default\"](dev))\n",
    "    # 设置输入数据\n",
    "    rt_mod.set_input(input_data=in_data)\n",
    "    # 设置模型参数\n",
    "    # rt_mod.set_input(**params)\n",
    "    # 运行模型\n",
    "    rt_mod.run()\n",
    "    # 获取输出结果\n",
    "    res = rt_mod.get_output(0).numpy()\n",
    "    return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7fd166ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_simulated_quantize(input_data, scale, zp, dtype, axis=-1):\n",
    "    \"\"\"\n",
    "    构建用于QNN模拟量化的虚拟机执行环境。\n",
    "\n",
    "    该函数创建一个包含QNN模拟量化操作的计算图，编译并优化它，\n",
    "    然后返回一个可以执行此量化操作的虚拟机实例。主要用于测试\n",
    "    模拟量化功能的性能和正确性。\n",
    "\n",
    "    参数:\n",
    "    input_data (relay.Expr): 输入的Relay表达式\n",
    "    scale (relay.Expr): 量化缩放因子\n",
    "    zp (relay.Expr): 量化零点\n",
    "    dtype (str): 输出数据类型，如 'int8', 'uint8'\n",
    "    axis (int, optional): 量化操作应用的轴，默认为-1（最后一个维度）\n",
    "\n",
    "    返回:\n",
    "    tvm.relay.vm.VirtualMachine: 可以执行模拟量化操作的虚拟机实例\n",
    "    \"\"\"\n",
    "    # 创建QNN模拟量化操作\n",
    "    sim_q = relay.qnn.simulated_quantize(\n",
    "        input_data,\n",
    "        scale,\n",
    "        zp,\n",
    "        axis=axis,\n",
    "        out_dtype=dtype,\n",
    "    )\n",
    "    # 从表达式构建IR模块\n",
    "    mod = tvm.IRModule.from_expr(sim_q)\n",
    "    # 编译模块，设置优化级别为3\n",
    "    with tvm.transform.PassContext(opt_level=3):\n",
    "        vm_exec = relay.vm.compile(mod, \"llvm\", params=None)\n",
    "    # 创建虚拟机实例\n",
    "    vm = VirtualMachine(vm_exec, tvm.cpu(0))\n",
    "    return vm\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a48020b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def verify_simulated_quantize_simple(dtype):\n",
    "    \"\"\"\n",
    "    验证QNN模拟量化功能的简单测试函数。\n",
    "\n",
    "    该函数通过生成随机测试数据，分别使用常规量化和模拟量化两种方式\n",
    "    执行量化操作，然后比较两种方式的结果是否一致，以验证模拟量化\n",
    "    功能的正确性。\n",
    "\n",
    "    参数:\n",
    "    dtype (str): 输出数据类型，如 'int8', 'uint8'\n",
    "\n",
    "    步骤:\n",
    "    1. 生成随机测试数据\n",
    "    2. 设置量化参数（缩放因子、零点）\n",
    "    3. 使用quantize_test_driver执行常规量化\n",
    "    4. 创建Relay变量和计算图\n",
    "    5. 构建模拟量化虚拟机\n",
    "    6. 执行模拟量化\n",
    "    7. 比较两种量化结果是否一致\n",
    "    \"\"\"\n",
    "    # 生成随机测试数据，范围[-128, 127)，形状[2, 5]，类型float32\n",
    "    data = np.random.uniform(low=-128, high=127, size=[2, 5]).astype(\"float32\")\n",
    "    # 设置量化缩放因子为0.5\n",
    "    scale_np = np.float32(0.5)\n",
    "    # 设置量化零点为127\n",
    "    zp_np = np.int32(127)\n",
    "    # 将数据类型字符串转换为对应的整数编码\n",
    "    dtype_np = np.int32(SQNN_DTYPE_TO_CODE[dtype])\n",
    "    # 准备量化参数字典\n",
    "    quant_args = {\"out_zero_point\": zp_np, \"out_scale\": scale_np}\n",
    "    # 使用常规量化方法得到预期输出\n",
    "    q_out = quantize_test_driver(\n",
    "        in_dtype=\"float32\",\n",
    "        quant_args=quant_args,\n",
    "        axis=-1,\n",
    "        out_dtype=dtype,\n",
    "        in_data=data,\n",
    "    )\n",
    "    # 创建Relay变量表示输入数据\n",
    "    input_data = relay.var(\"input_data\", shape=data.shape, dtype=\"float32\")\n",
    "    # 创建Relay变量表示缩放因子\n",
    "    scale = relay.var(\"scale\", shape=[])\n",
    "    # 创建Relay变量表示零点\n",
    "    zp = relay.var(\"zp\", shape=[], dtype=\"int32\")\n",
    "    # 创建Relay变量表示数据类型\n",
    "    dtype = relay.var(\"dtype\", shape=[], dtype=\"int32\")\n",
    "    # 构建模拟量化虚拟机\n",
    "    vm = build_simulated_quantize(input_data, scale, zp, dtype)\n",
    "    # 在虚拟机上执行模拟量化\n",
    "    sim_q_out = vm.invoke(\"main\", input_data=data, scale=scale_np, zp=zp_np, dtype=dtype_np)\n",
    "    # 比较模拟量化结果与预期输出是否一致（容忍浮点误差）\n",
    "    allclose_with_rounding(sim_q_out.numpy(), q_out)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97f94ff2",
   "metadata": {},
   "source": [
    "## QNN 模拟量化功能测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9055b45a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_simulated_quantize():\n",
    "    \"\"\"\n",
    "    QNN模拟量化功能的主测试函数。\n",
    "\n",
    "    该函数通过调用verify_simulated_quantize_simple函数，对三种不同的\n",
    "    输出数据类型（无符号8位整数、有符号8位整数和有符号32位整数）\n",
    "    进行模拟量化测试，以验证模拟量化功能在不同数据类型下的正确性。\n",
    "    \"\"\"\n",
    "    # 测试无符号8位整数（uint8）的模拟量化\n",
    "    verify_simulated_quantize_simple(\"uint8\")\n",
    "    # 测试有符号8位整数（int8）的模拟量化\n",
    "    verify_simulated_quantize_simple(\"int8\")\n",
    "    # 测试有符号32位整数（int32）的模拟量化\n",
    "    verify_simulated_quantize_simple(\"int32\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ccb87db0",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_simulated_quantize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9748836",
   "metadata": {},
   "source": [
    "## 测试动态通道的QNN模拟量化功能"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "384e4073",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_dynamic_channels():\n",
    "    \"\"\"\n",
    "    测试动态通道的QNN模拟量化功能。\n",
    "    \n",
    "    该函数验证模拟量化能否在不重新编译的情况下，同时支持标量量化参数和每通道量化参数。\n",
    "    主要测试流程：\n",
    "    1. 生成随机测试数据\n",
    "    2. 使用标量量化参数进行测试\n",
    "    3. 创建具有未定义形状的变量并构建模拟量化虚拟机\n",
    "    4. 使用标量参数运行模拟量化并验证结果\n",
    "    5. 使用每通道量化参数进行测试（无需重新编译虚拟机）\n",
    "    6. 验证每通道量化结果\n",
    "    \n",
    "    此测试重点在于验证模拟量化对动态通道配置的支持能力。\n",
    "    \"\"\"\n",
    "    # 编译一次模拟量化，但支持每通道或标量参数\n",
    "    data = np.random.uniform(low=-64, high=64, size=[2, 5]).astype(\"float32\")\n",
    "    \n",
    "    # 测试标量QNN参数\n",
    "    scale_np = np.asarray([0.5]).astype(\"float32\")  # 标量缩放因子\n",
    "    zp_np = np.asarray([127]).astype(\"int32\")       # 标量零点\n",
    "    dtype_np = np.int32(SQNN_DTYPE_TO_CODE[\"uint8\"])  # 输出数据类型代码\n",
    "    quant_args = {\"out_zero_point\": zp_np[0], \"out_scale\": scale_np[0]}\n",
    "    \n",
    "    # 使用测试驱动函数获取量化输出\n",
    "    q_out = quantize_test_driver(\n",
    "        in_dtype=\"float32\",    # 输入数据类型\n",
    "        quant_args=quant_args, # 量化参数\n",
    "        axis=0,                # 量化轴\n",
    "        out_dtype=\"uint8\",     # 输出数据类型\n",
    "        in_data=data,          # 输入数据\n",
    "    )\n",
    "    \n",
    "    # 创建具有未定义形状的变量，以支持动态参数\n",
    "    input_data = relay.var(\"input_data\", shape=data.shape, dtype=\"float32\")\n",
    "    scale = relay.var(\"scale\", shape=[relay.Any()], dtype=\"float32\")  # 使用relay.Any()允许任意形状\n",
    "    zp = relay.var(\"zp\", shape=[relay.Any()], dtype=\"int32\")\n",
    "    dtype = relay.var(\"dtype\", shape=[], dtype=\"int32\")\n",
    "    \n",
    "    # 构建模拟量化虚拟机\n",
    "    vm = build_simulated_quantize(input_data, scale, zp, dtype, axis=0)\n",
    "    \n",
    "    # 使用标量输入运行模拟量化\n",
    "    sim_q_out = vm.invoke(\"main\", input_data=data, scale=scale_np, zp=zp_np, dtype=dtype_np)\n",
    "    \n",
    "    # 验证模拟量化结果与预期一致\n",
    "    allclose_with_rounding(sim_q_out.numpy(), q_out)\n",
    "\n",
    "    # 测试每通道量化参数，无需重新编译虚拟机\n",
    "    scale_np = np.array([0.5, 0.25]).astype(\"float32\")  # 每通道缩放因子\n",
    "    zp_np = np.array([127, 123]).astype(\"int32\")        # 每通道零点\n",
    "\n",
    "    # 获取参考量化输出\n",
    "    quant_args = {\"out_zero_point\": zp_np, \"out_scale\": scale_np}\n",
    "    q_out = quantize_test_driver(\n",
    "        in_dtype=\"float32\",\n",
    "        quant_args=quant_args,\n",
    "        axis=0,\n",
    "        out_dtype=\"uint8\",\n",
    "        in_data=data,\n",
    "    )\n",
    "    \n",
    "    # 在不重新编译的情况下运行模拟量化并验证结果\n",
    "    sim_q_out = vm.invoke(\"main\", input_data=data, scale=scale_np, zp=zp_np, dtype=dtype_np)\n",
    "    allclose_with_rounding(sim_q_out.numpy(), q_out)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "37ca89a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dynamic_channels()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7aedb802",
   "metadata": {},
   "source": [
    "## 测试动态数据类型的QNN模拟量化功能"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ddbe8a6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_dynamic_dtype():\n",
    "    \"\"\"\n",
    "    测试动态数据类型的QNN模拟量化功能。\n",
    "    \n",
    "    该函数验证模拟量化能否在不重新编译的情况下，支持不同的输出数据类型量化。\n",
    "    主要测试流程：\n",
    "    1. 生成随机测试数据\n",
    "    2. 测试float32到uint8的标量量化\n",
    "    3. 创建具有未定义形状的变量并构建模拟量化虚拟机\n",
    "    4. 使用标量参数运行模拟量化并验证结果\n",
    "    5. 测试float32到int32的量化（无需重新编译虚拟机）\n",
    "    6. 验证int32量化结果\n",
    "    \n",
    "    此测试重点在于验证模拟量化对动态数据类型配置的支持能力。\n",
    "    \"\"\"\n",
    "    # 编译一次模拟量化，但支持不同类型的量化\n",
    "    data = np.random.uniform(low=-64, high=64, size=[2, 5]).astype(\"float32\")\n",
    "    \n",
    "    # 测试float32到uint8的标量量化\n",
    "    scale_np = np.asarray([0.5]).astype(\"float32\")  # 标量缩放因子\n",
    "    zp_np = np.asarray([127]).astype(\"int32\")       # 标量零点\n",
    "    dtype_np = np.int32(SQNN_DTYPE_TO_CODE[\"uint8\"])  # uint8类型代码\n",
    "    quant_args = {\"out_zero_point\": zp_np[0], \"out_scale\": scale_np[0]}\n",
    "    \n",
    "    # 使用测试驱动函数获取uint8量化输出\n",
    "    q_out = quantize_test_driver(\n",
    "        in_dtype=\"float32\",    # 输入数据类型\n",
    "        quant_args=quant_args, # 量化参数\n",
    "        axis=-1,               # 量化轴（最后一维）\n",
    "        out_dtype=\"uint8\",     # 输出数据类型\n",
    "        in_data=data,          # 输入数据\n",
    "    )\n",
    "    \n",
    "    # 创建具有未定义形状的变量，以支持动态参数\n",
    "    input_data = relay.var(\"input_data\", shape=data.shape, dtype=\"float32\")\n",
    "    scale = relay.var(\"scale\", shape=[relay.Any()], dtype=\"float32\")  # 任意形状的缩放因子\n",
    "    zp = relay.var(\"zp\", shape=[relay.Any()], dtype=\"int32\")         # 任意形状的零点\n",
    "    dtype = relay.var(\"dtype\", shape=[], dtype=\"int32\")              # 数据类型变量\n",
    "    \n",
    "    # 构建模拟量化虚拟机\n",
    "    vm = build_simulated_quantize(input_data, scale, zp, dtype)\n",
    "    \n",
    "    # 使用标量输入运行模拟量化（uint8）\n",
    "    sim_q_out = vm.invoke(\"main\", input_data=data, scale=scale_np, zp=zp_np, dtype=dtype_np)\n",
    "    \n",
    "    # 验证模拟量化结果与预期一致\n",
    "    allclose_with_rounding(sim_q_out.numpy(), q_out)\n",
    "\n",
    "    # 测试float32到int32的量化，无需重新编译虚拟机\n",
    "    \n",
    "    # 获取int32参考量化输出\n",
    "    q_out = quantize_test_driver(\n",
    "        in_dtype=\"float32\",\n",
    "        quant_args=quant_args,\n",
    "        axis=-1,\n",
    "        out_dtype=\"int32\",\n",
    "        in_data=data,\n",
    "    )\n",
    "    \n",
    "    # 更新数据类型为int32\n",
    "    dtype_np = np.int32(SQNN_DTYPE_TO_CODE[\"int32\"])\n",
    "    \n",
    "    # 在不重新编译的情况下运行模拟量化（int32）并验证结果\n",
    "    sim_q_out = vm.invoke(\"main\", input_data=data, scale=scale_np, zp=zp_np, dtype=dtype_np)\n",
    "    allclose_with_rounding(sim_q_out.numpy(), q_out)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "8d972102",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dynamic_dtype()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b0e7444",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "aix",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
